--[[ 

 This module provides lusi's gillespie simulation driver

 (c) 2013 Markus Dittrich

 This program is free software; you can redistribute it
 and/or modify it under the terms of the GNU General Public
 License Version 3 as published by the Free Software Foundation.

 This program is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU General Public License Version 3 for more details.

 You should have received a copy of the GNU General Public
 License along with this program; if not, write to the Free
 Software Foundation, Inc., 59 Temple Place - Suite 330,
 Boston, MA 02111-1307, USA.

--]]

-- use this only during development
-- tostring = require 'ml'.tstring

-- load modules
local random = require "random"
local math = require "math"
local rpn_engine = require "rpn_engine"


-- basic defitions
local volume = 5e-20         
local N_avo  = 6.0221415e23


-- computes the reaction proensities given the current
-- molecule counts
function compute_propensities(molecules, reactions, variables) 

  local props = {}
  for _, reaction in ipairs(reactions) do

    local hVal = 1
    for _, reactant in ipairs(reaction.reacts) do
      local h = 1
      for i=1,reactant.mult do
        h = h * (reactant.mol.count - i + 1)
      end
      h = 1/reactant.mult * h
      hVal = hVal * h
    end

    -- determine the current rate, may require evaluation of an 
    -- RPN expression
    local rate = reaction.rate
    if type(rate) == "table" then
      rate = rpn_engine.compute_rpn(reaction.rate, variables, molecules)
      assert(rate, "Error: failed to compute reaction rate")
    end
    props[#props+1] = hVal * rate
  end

  return props
end



-- update the molecule count based on the pathway chosen
function update_populations(reactions, mu)

  local reaction = reactions[mu]
  for _,reactant in ipairs(reaction.reacts) do
    reactant.mol.count = reactant.mol.count - reactant.mult
  end

  for _, prod in ipairs(reaction.products) do
    prod.mol.count = prod.mol.count + prod.mult
  end

end



-- return the sum of all values in the array
-- NOTE: the array is assumed to be numeric
function sum_array(array) 

  local sum = 0
  for _, val in ipairs(array) do
    sum = sum + val
  end

  return sum
end


-- writes the current molecule counts to file
function write_output(outfiles, molecules, time) 
  for name, outfile in pairs(outfiles) do
    outfile:write(string.format("%f  %d\n", time, molecules[name].count))
  end

end


-- close files
function close_files()
  for _, outfile in ipairs(outFiles) do
    outfile:close()
  end
end


-- open output files
function open_files(molecules)
  local outFiles = {}
  for name, mol in pairs(molecules) do
    outFiles[name] = io.open(mol.name .. ".dat", "w")
  end
  return outFiles
end


-- main loop
function run_gillespie(molecules, reactions, params, variables) 

  local time = 0
  local itercount = 1
  local r = random.uniform()
  local length = params.time
  local outputFreq = params.interval

  outFiles = open_files(molecules)

  while time < length do

    -- compute propensities
    local a_i = compute_propensities(molecules, reactions, variables)
    local a_0 = sum_array(a_i)

    -- make sure reaction hasn't stopped
    if (a_0 - 0) < 1e-16 then
      print("WARNING: reaction has stopped")
      break;
    end

    -- compute tau
    local tau = (1/a_0) * math.log(1/r())
    time = time + tau

    -- compute mu
    local mu = 0
    local thresh = r() * a_0
    local sum = 0
    while sum < thresh do
      sum = sum + a_i[mu+1]
      mu = mu + 1
    end

    update_populations(reactions, mu)  

    if (itercount % outputFreq) == 0 then
      write_output(outFiles, molecules, time)
      print(string.format("%d:  time = %f", itercount, time))
    end 
  
    itercount = itercount + 1
  end

  close_files()
end


gillespie_engine = {
    run_gillespie = run_gillespie 
}

return gillespie_engine
